import os
import h5py
import cv2

def change_hdf5_format(dataset_path,output_path):
    with h5py.File(dataset_path, 'r') as root:
        action = root['/action'][()]
        qpos = root['/observations/qpos'][()]
        base_action = root['/base_action'][()]
        compressed = root.attrs.get('compress', False)
        image_dict = dict()
        for cam_name in root[f'/observations/images/'].keys():
            image_dict[cam_name] = root[f'/observations/images/{cam_name}'][()]
    
    with h5py.File(output_path, 'w') as f:
        f.attrs['compress'] = compressed
        # action从(300,14)裁剪为(300,7)
        action = action[:, :7]
        f.create_dataset('action', data=action)
        # base_action
        f.create_dataset('base_action', data=base_action)

        f.create_group('observations')
        # qpos从(300,14)裁剪为(300,7)
        qpos = qpos[:, :7]
        f['observations'].create_dataset('qpos', data=qpos)

        # images原本有cam_high,cam_left_wrist,cam_right_wrist，舍弃cam_right_wrist
        image_dict = {k: v for k, v in image_dict.items() if k != 'cam_right_wrist'}
        f['observations'].create_group('images')
        for cam_name, image_list in image_dict.items():
            f['observations/images'].create_dataset(cam_name, data=image_list)

if __name__ == "__main__":
    #读取/gemini/data-1/ACT/yogurt_grap_compressed 下的全部hdf5文件，并修改为/gemini/data-1/ACT/yogurt_new 下的hdf5文件
    dataset_dir = "/gemini/data-1/ACT/yogurt_compressed"
    output_dir = "/gemini/data-1/ACT/yogurt_single_arm"
    os.makedirs(output_dir, exist_ok=True)
    for file in os.listdir(dataset_dir):
        if file.endswith(".hdf5"):
            dataset_path = os.path.join(dataset_dir, file)
            output_path = os.path.join(output_dir, file)
            change_hdf5_format(dataset_path, output_path)